import json
import logging
import os
import random
import threading
import time
import collections
from typing import Optional

from flask import Flask, request, abort
from nacos_py import NacosService, NacosClient
from dotenv import load_dotenv
from arch_inference import inference as _inference



IP = "localhost"
PORT = "12488"

NACOS_ENABLE = True

SERVER_ADDRESSES = "http://localhost:8848"
NAMESPACE = "public"
CLIENT: Optional[NacosClient] = None
SERVICE_NAME = "nlp"

NAME = "nacos"  # Nacos用户名
PWD = "nacos"  # Nacos密码

# Nacos 配置管理
config = {
    "GITHUB_TOKEN": "",
    "OPENAI_API_KEY": "",
    "HTTP_PROXY": "",
    "HTTPS_PROXY": "",
    "ALL_PROXY": ""
}

def heartbeat():
    while True:
        CLIENT.send_heartbeat("DEFAULT_GROUP@@" + SERVICE_NAME, IP, PORT)
        time.sleep(3)

logging.basicConfig(level=logging.DEBUG)
load_dotenv()
if NACOS_ENABLE:
    CLIENT = NacosClient(SERVER_ADDRESSES, namespace=NAMESPACE, username=NAME, password=PWD)
    NacosService(SERVICE_NAME, CLIENT).register(IP, PORT)
    heart = threading.Thread(target=heartbeat, daemon=True)
    heart.start()


class Inference(Flask):
    def __init__(self, *args, **kwargs):
        super(Inference, self).__init__(*args, **kwargs)


app = Inference(__name__)




def get_service(service_name):
    instances = CLIENT.list_naming_instance(service_name, namespace_id=NAMESPACE, healthy_only=True)
    if len(instances["hosts"][0]) == 0:
        raise Exception(f"服务调用异常:没有找到指定服务-{service_name}")
    # 负载均衡
    host = random.choice(instances["hosts"])
    return host["ip"], host["port"]

def updated_config():
    global config
    data_id = "inference.json"
    config = json.loads(CLIENT.get_config(data_id, "DEFAULT_GROUP", no_snapshot=True))
    for _config in config.keys():
        os.environ[_config] = config[_config]


@app.route("/inference", methods=["POST"])
def inference():
    try:
        raw_data = request.get_data(as_text=True)
        commits = json.loads(raw_data)
    except:
        abort(400)
    repo_commits = collections.defaultdict(lambda: [])
    repo_index = collections.defaultdict(lambda: [])
    for idx in range(len(commits)):
        repo_commits[commits[idx]["repo_id"]].append(commits[idx]["hash"])
        repo_index[commits[idx]["repo_id"]].append(idx)
    result = [{}] * len(commits)
    for repo_id in repo_commits.keys():
        _results = _inference(repo_id, repo_commits[repo_id])
        for idx in range(len(_results)):
            result[repo_index[repo_id][idx]] = {
                "result": _results[idx]
            }
    return result

if __name__ == '__main__':

    app.run(host=IP, port=int(PORT), debug=True)
